import json
import os
import tempfile
import zipfile
from pathlib import Path
from typing import List, Dict, Any, Optional

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from fastapi import APIRouter, UploadFile, File, Form, HTTPException

router = APIRouter()

IMAGE_EXTENSIONS = {"jpg", "jpeg", "png", "gif", "bmp", "svg", "webp"}
MAX_FILE_BYTES = 100 * 1024 * 1024  # 100 MB

# 模型配置
MODEL_PATH = os.getenv("AUDIT_MODEL_PATH", "/data/vlm/mzh/finetune-merge-deepseek-8b")
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"


class CodeAuditModel:
    """代码审计模型加载器"""
    _instance = None
    _model = None
    _tokenizer = None
    
    def __new__(cls):
        if cls._instance is None:
            cls._instance = super().__new__(cls)
        return cls._instance
    
    def load_model(self):
        """加载模型(单例模式)"""
        if self._model is None:
            print(f"Loading model from {MODEL_PATH}...")
            self._tokenizer = AutoTokenizer.from_pretrained(
                MODEL_PATH,
                trust_remote_code=True
            )
            self._model = AutoModelForCausalLM.from_pretrained(
                MODEL_PATH,
                torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32,
                device_map="auto" if DEVICE == "cuda" else None,
                trust_remote_code=True
            )
            if DEVICE == "cpu":
                self._model = self._model.to(DEVICE)
            self._model.eval()
            print(f"Model loaded successfully on {DEVICE}")
        return self._model, self._tokenizer
    
    def generate(self, prompt: str, code: str, max_new_tokens: int = 40, temperature: float = 0.7, top_p: float = 0.95) -> str:
        """生成文本"""
        model, tokenizer = self.load_model()
        
        # 使用原模型的chat模板
        messages = [
            {"role": "system", "content": prompt},
            {"role": "user", "content": code},
        ]
        
        # 使用tokenizer的chat模板
        inputs = tokenizer.apply_chat_template(
            messages,
            tokenize=True,
            add_generation_prompt=True,
            return_dict=True,
            return_tensors="pt"
        ).to(model.device)
        print(f"[DEBUG] 使用原模型chat模板")

        
        
        input_length = inputs.input_ids.shape[1]
        print(f"[DEBUG] 输入token数量: {input_length}")
        
        outputs = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            temperature=temperature,
            top_p=top_p,
        )
        
        output_length = outputs[0].shape[0]
        generated_tokens: Any = output_length - input_length
        print(f"[DEBUG] 输出token总数: {output_length}")
        print(f"[DEBUG] 生成的token数: {generated_tokens}")
        
        generated_text = tokenizer.decode(outputs[0][inputs["input_ids"].shape[-1]:], skip_special_tokens=True)
        
        
        '''# 清理可能的标签
        generated_text: Any = generated_text.strip()
        
        # 移除常见的模型输出标签
        for tag in ["<|system|>", "<|user|>", "<|assistant|>", "<|end|>", "<think>", "</think>"]:
            if tag in generated_text:
                print(f"[DEBUG] 移除标签: {tag}")
            generated_text = generated_text.replace(tag, "")
        
        final_text = generated_text.strip()'''
        print(f"[DEBUG] 最终输出长度: {len(generated_text)} 字符")
        
        return generated_text


# 全局模型实例
audit_model = CodeAuditModel()


def sanitize_file_name(filename: str) -> str:
    safe = "".join(ch if ch.isalnum() or ch in {".", "_", "-", "/"} else "_" for ch in filename)
    return safe[:255] or "unnamed"


def is_image(path: Path) -> bool:
    return path.suffix.lower().lstrip(".") in IMAGE_EXTENSIONS


def read_text(path: Path) -> str:
    try:
        return path.read_text(encoding="utf-8")
    except UnicodeDecodeError:
        return path.read_text(encoding="utf-8", errors="ignore")


def extract_zip_to_sources(zip_path: Path) -> List[Dict[str, Any]]:
    sources: List[Dict[str, Any]] = []
    with zipfile.ZipFile(zip_path) as zf:
        tmp_root = Path(tempfile.mkdtemp(prefix="audit_zip_"))
        zf.extractall(tmp_root)
        for entry in tmp_root.rglob("*"):
            if not entry.is_file():
                continue
            if is_image(entry):
                continue
            if entry.stat().st_size > MAX_FILE_BYTES:
                continue
            rel = entry.relative_to(tmp_root)
            sources.append({
                "path": sanitize_file_name(str(rel)),
                "code": read_text(entry)
            })
    return sources


def generate_prompt(language: str, categories: List[str], items: List[str]) -> str:
    # 使用精简的prompt，要求模型只输出CWE代码列表
    return (
        "You are an expert in the field of code. Please determine the CWE types present in the following code snippet.\n"
        "- If there are no vulnerabilities in the code, output an empty string (\"\").\n"
        "- If multiple vulnerabilities exist, the output format should be: \"CWE-xxx,CWE-yyy\" (within double quotes, no spaces, sorted in ascending order by CWE number).\n"
        "- Return only the CWE numbers, without any explanations.\n"
        "- If the same type occurs multiple times, output it only once (e.g., if there are two instances of CWE-401, output only \"CWE-401\").\n"
        "\n"
        "Example:\n"
        "Input: [Vulnerable Code Snippet]\n"
        "Output: \"CWE-401,CWE-402\"\n"
    )


def call_audit_model(code: str, language: str, categories: List[str], items: List[str], model: str) -> List[Dict[str, Any]]:
    if not code.strip():
        return []

    prompt = generate_prompt(language, categories, items)
    
    try:
        # 使用本地模型生成
        content = audit_model.generate(
            prompt,
            code,
            temperature=0.6,
            top_p=0.95
        )
        print(f"\n模型生成内容:\n{content}")
        print(f"\n模型生成内容总长度: {len(content)} 字符")
    except Exception as e:
        print(f"Model generation error: {e}")
        return []

    cwes = extract_cwe_codes(content)
    print(f"提取到的CWE代码: {cwes}")
    
    if not cwes:
        print("未提取到CWE代码，视为无安全风险")
        # 未提取到CWE代码，返回空列表表示无风险
        return []

    severity = aggregate_severity(cwes)
    description = " | ".join(cwe_description(c) for c in cwes)

    return [{
        "vulnerabilityType": "/".join(cwes),
        "lineNumber": 0,
        "severity": severity,
        "description": description,
        "codeSnippet": code,
        "payload": "",
    }]


def extract_cwe_codes(text: str) -> List[str]:
    if not text:
        return []
    text = text.split("</think>")[-1]
    matches = []
    seen = set()
    for part in text.split():
        if part.upper().startswith("CWE-"):
            code = part.upper().rstrip("|).,:")
            if code not in seen:
                seen.add(code)
                matches.append(code)
    return matches


SEVERITY_RANK = {"critical": 4, "high": 3, "medium": 2, "low": 1, "info": 0}


def aggregate_severity(cwes: List[str]) -> str:
    best = "medium"
    best_rank = SEVERITY_RANK[best]
    for c in cwes:
        sev = cwe_severity(c)
        rank = SEVERITY_RANK.get(sev, 2)
        if rank > best_rank:
            best, best_rank = sev, rank
            if best == "critical":
                break
    return best


CWE_DESCRIPTION: Dict[str, str] = {
    'CWE-16': '系统安全配置存在缺陷，可能被攻击者利用',
    'CWE-20': '没有正确验证用户输入，可能导致注入攻击或内存错误',
    'CWE-22': '路径遍历漏洞，攻击者可以访问系统任意文件',
    'CWE-73': '文件路径受外部控制，可能被重定向到其他位置',
    'CWE-75': '特殊字符未被正确处理，可能被用于命令或查询注入',
    'CWE-77': '命令注入漏洞，攻击者可以执行系统命令',
    'CWE-78': '操作系统命令注入，可完全控制服务器',
    'CWE-79': '跨站脚本攻击(XSS)，可在用户浏览器中执行恶意脚本',
    'CWE-80': 'XSS漏洞，未正确过滤用户输入的脚本代码',
    'CWE-89': 'SQL注入漏洞，可篡改或窃取数据库信息',
    'CWE-93': 'CRLF注入，可能导致HTTP响应拆分攻击',
    'CWE-94': '代码注入漏洞，攻击者可执行任意代码',
    'CWE-95': '动态执行用户输入，可能导致代码注入',
    'CWE-102': '框架配置不当，存在安全隐患',
    'CWE-113': 'HTTP响应头注入，可篡改HTTP响应',
    'CWE-116': '输出编码不当，容易被注入攻击',
    'CWE-117': '日志记录未过滤，可能被伪造或泄露信息',
    'CWE-119': '缓冲区边界检查不严，可能导致溢出',
    'CWE-120': '栈缓冲区溢出，可被利用执行恶意代码',
    'CWE-125': '读取了不该读的内存区域，可能泄露敏感信息',
    'CWE-134': '格式化字符串漏洞，可导致程序崩溃或代码执行',
    'CWE-190': '整数运算溢出，可能导致逻辑错误或安全问题',
    'CWE-200': '敏感信息意外泄露给未授权用户',
    'CWE-201': '通过网络传输泄露敏感数据',
    'CWE-209': '错误信息中包含敏感内容，可能被攻击者利用',
    'CWE-255': '密码或密钥管理存在安全风险',
    'CWE-256': '密码以明文形式存储，极易被窃取',
    'CWE-259': '在代码中直接写死密码，非常不安全',
    'CWE-269': '权限管理混乱，可能导致越权操作',
    'CWE-284': '访问控制不严，未授权用户可访问资源',
    'CWE-285': '权限检查不充分，可能导致越权访问',
    'CWE-287': '身份认证机制有缺陷，可能被绕过',
    'CWE-295': '证书验证不严格，可能遭受中间人攻击',
    'CWE-306': '关键功能缺少身份验证，任何人都可访问',
    'CWE-319': '敏感数据在网络中明文传输',
    'CWE-321': '在代码中硬编码加密密钥，极易被发现',
    'CWE-322': '密钥交换时未验证对方身份',
    'CWE-326': '加密算法强度不够，容易被破解',
    'CWE-327': '使用了已知的弱加密算法',
    'CWE-328': '使用了安全性差的哈希算法',
    'CWE-330': '随机数生成质量差，容易被预测',
    'CWE-331': '随机数生成器熵值不足',
    'CWE-338': '使用了密码学上不安全的随机数生成器',
    'CWE-352': '跨站请求伪造，可诱使用户执行非预期操作',
    'CWE-359': '用户隐私数据处理不当，可能被泄露',
    'CWE-384': '会话固定漏洞，可能导致会话劫持',
    'CWE-398': '代码质量差，存在潜在安全风险',
    'CWE-399': '资源管理不当，可能导致系统异常',
    'CWE-400': '资源耗尽攻击，可能导致服务拒绝',
    'CWE-404': '资源未正确释放，可能导致内存泄漏',
    'CWE-416': '使用了已释放的内存，可能导致程序崩溃',
    'CWE-428': '搜索路径未加引号，可能执行恶意程序',
    'CWE-434': '文件上传限制不严，可能导致远程执行',
    'CWE-476': '访问了空指针，可能导致程序崩溃',
    'CWE-502': '反序列化不可信数据，可导致远程代码执行',
    'CWE-511': '逻辑或时序问题，可能引发安全漏洞',
    'CWE-522': '用户凭据保护不充分',
    'CWE-532': '敏感信息被记录在日志中',
    'CWE-568': '对象销毁处理不当，可能导致异常行为',
    'CWE-611': 'XML外部实体注入，可读取服务器文件',
    'CWE-614': 'HTTPS下Cookie缺少安全属性',
    'CWE-639': '直接对象引用未验证，可越权访问数据',
    'CWE-668': '资源暴露给了不应该访问的用户',
    'CWE-676': '使用了危险的系统函数',
    'CWE-732': '关键文件或目录权限设置错误',
    'CWE-749': '暴露了不安全的函数或方法',
    'CWE-759': '密码哈希未加盐，容易被破解',
    'CWE-772': '资源使用完后未及时释放',
    'CWE-776': 'XML实体递归引用，可能导致拒绝服务',
    'CWE-780': 'RSA加密未使用OAEP填充',
    'CWE-787': '向内存写入了不该写的数据，可能导致溢出',
    'CWE-798': '在代码中硬编码用户名密码',
    'CWE-835': '死循环无法退出，可能导致系统卡死',
    'CWE-917': '表达式语言注入，可执行任意代码',
    'CWE-918': '服务端请求伪造，可探测内网服务',
    'CWE-919': '移动端界面存在安全缺陷，可能被钓鱼',
    'CWE-1035': '认证机制太弱，容易被暴力破解',
    'CWE-1204': '通信端点限制不严，可能被恶意连接',
}


def cwe_description(code: str) -> str:
    return CWE_DESCRIPTION.get(code.upper(), f"{code.upper()} vulnerability")


CWE_SEVERITY: Dict[str, str] = {
    # Critical
    'CWE-22': 'critical', 'CWE-73': 'critical', 'CWE-78': 'critical', 'CWE-89': 'critical', 'CWE-94': 'critical',
    'CWE-119': 'critical', 'CWE-120': 'critical', 'CWE-125': 'critical', 'CWE-134': 'critical', 'CWE-306': 'critical',
    'CWE-322': 'critical', 'CWE-434': 'critical', 'CWE-502': 'critical', 'CWE-787': 'critical',
    # High
    'CWE-16': 'high', 'CWE-20': 'high', 'CWE-77': 'high', 'CWE-79': 'high', 'CWE-80': 'high', 'CWE-113': 'high',
    'CWE-116': 'high', 'CWE-190': 'high', 'CWE-209': 'high', 'CWE-255': 'high', 'CWE-269': 'high', 'CWE-284': 'high',
    'CWE-285': 'high', 'CWE-287': 'high', 'CWE-295': 'high', 'CWE-319': 'high', 'CWE-321': 'high', 'CWE-326': 'high',
    'CWE-330': 'high', 'CWE-331': 'high', 'CWE-338': 'high', 'CWE-352': 'high', 'CWE-359': 'high', 'CWE-384': 'high',
    'CWE-399': 'high', 'CWE-428': 'high', 'CWE-476': 'high', 'CWE-522': 'high', 'CWE-611': 'high', 'CWE-614': 'high',
    'CWE-639': 'high', 'CWE-668': 'high', 'CWE-732': 'high', 'CWE-749': 'high', 'CWE-772': 'high', 'CWE-776': 'high',
    'CWE-798': 'high', 'CWE-835': 'high', 'CWE-917': 'high', 'CWE-918': 'high', 'CWE-919': 'high', 'CWE-1204': 'high',
    # Medium
    'CWE-75': 'medium', 'CWE-93': 'medium', 'CWE-95': 'medium', 'CWE-102': 'medium', 'CWE-200': 'medium',
    'CWE-201': 'medium', 'CWE-256': 'medium', 'CWE-259': 'medium', 'CWE-327': 'medium', 'CWE-328': 'medium',
    'CWE-398': 'medium', 'CWE-400': 'medium', 'CWE-404': 'medium', 'CWE-416': 'medium', 'CWE-759': 'medium',
    'CWE-780': 'medium', 'CWE-1035': 'medium',
    # Low
    'CWE-117': 'low', 'CWE-532': 'low', 'CWE-563': 'low', 'CWE-676': 'low', 'CWE-1041': 'low',
}


def cwe_severity(code: str) -> str:
    return CWE_SEVERITY.get(code.upper(), "medium")


def compute_summary(issues: List[Dict[str, Any]]) -> Dict[str, int]:
    summary = {"critical": 0, "high": 0, "medium": 0, "low": 0, "info": 0}
    for issue in issues:
        sev = str(issue.get("severity", "medium")).lower()
        if sev not in summary:
            sev = "medium"
        summary[sev] += 1
    return summary


@router.post("/api/audit")
async def audit_endpoint(
    language: str = Form(""),
    categories: str = Form("[]"),
    items: str = Form("[]"),
    text_content: str = Form(""),
    model: str = Form("qwen2.5-coder-instruct"),
    file: Optional[UploadFile] = File(None),
):
    print("\n=== 开始代码审计 ===")
    print(f"language: {language}")
    print(f"categories: {categories}")
    print(f"items: {items}")
    print(f"text_content length: {len(text_content)}")
    print(f"model: {model}")
    print(f"file: {file.filename if file else None}")
    
    try:
        categories_list = json.loads(categories) if categories else []
    except json.JSONDecodeError:
        categories_list = []

    try:
        items_list = json.loads(items) if items else []
    except json.JSONDecodeError:
        items_list = []

    sources: List[Dict[str, str]] = []

    if file and file.filename:
        suffix = Path(file.filename).suffix.lower()
        with tempfile.NamedTemporaryFile("wb", delete=False, suffix=suffix) as tmp:
            content = await file.read()
            tmp.write(content)
            tmp_path = Path(tmp.name)
        if tmp_path.stat().st_size > MAX_FILE_BYTES:
            tmp_path.unlink(missing_ok=True)
            raise HTTPException(status_code=400, detail="Uploaded file too large")

        if suffix == ".zip":
            try:
                sources.extend(extract_zip_to_sources(tmp_path))
            finally:
                tmp_path.unlink(missing_ok=True)
        else:
            code = read_text(tmp_path)
            sources.append({
                "path": sanitize_file_name(file.filename),
                "code": code,
            })
            tmp_path.unlink(missing_ok=True)

    if text_content and text_content.strip():
        sources.append({"path": "text-input.txt", "code": text_content})
        print(f"\u6dfb加文本输入: {len(text_content)} 个字符")

    if not sources:
        print("错误: 没有提供文件或文本内容")
        raise HTTPException(status_code=400, detail="No file or text content provided")
    
    print(f"\u603b计 {len(sources)} 个源文件")

    all_issues: List[Dict[str, Any]] = []

    for idx, src in enumerate(sources):
        print(f"\n处理源文件 {idx + 1}/{len(sources)}: {src['path']}")
        print(f"代码长度: {len(src['code'])} 个字符")
        try:
            issues = call_audit_model(src["code"], language, categories_list, items_list, model)
            print(f"发现 {len(issues)} 个问题")
        except Exception as exc:  # noqa: BLE001
            raise HTTPException(status_code=500, detail=f"Model error: {str(exc)}") from exc

        for issue in issues:
            issue_copy = dict(issue)
            issue_copy.setdefault("file", src["path"])
            issue_copy.setdefault("line", issue_copy.pop("lineNumber", 0))
            issue_copy["severity"] = str(issue_copy.get("severity", "medium")).lower()
            all_issues.append(issue_copy)

    summary = compute_summary(all_issues)
    
    print(f"\n=== 审计完成 ===")
    print(f"总问题数: {len(all_issues)}")
    print(f"问题统计: {summary}")

    return {
        "success": True,
        "code": 200,
        "message": "Audit completed successfully",
        "data": {
            "issues": all_issues,
            "count": len(all_issues),
            "summary": summary,
        },
    }
